import sys,os
import numpy as np
from textgrid import TextGrid, IntervalTier
from multiprocessing import Pool
ctm = {}

def write_textgrid(out_dir, i):
    global ctm
    tg = TextGrid()
    tg.append(ctm[i])
    path = f'{out_dir}/{i}.TextGrid'
    tg.write(path)

    return i, os.path.abspath(path)


def ctm_to_textgrid(ctm_file, out_dir):
    os.makedirs(out_dir, exist_ok=True)
    global ctm
    n=0
    for i in open(ctm_file):
        tmp = i.strip().split()
        if tmp[0] not in ctm:
            n += 1
            ctm[tmp[0]] = IntervalTier()
            ctm[tmp[0]].name = 'word'
        ctm[tmp[0]].add(float(tmp[2]), float(format(float(tmp[2]) + float(tmp[3]), '.3f')), tmp[4])
    
        #np.savez_compressed(f'{out_dir}/Textgrid', **ctm)

    print (n)
    p = Pool(30)

    f = open(f'{out_dir}/../textgrid.scp', 'w')
    res = []
    for i in ctm.keys():
        res.append(p.apply_async(write_textgrid, args=[out_dir, i]))
        # tg = TextGrid()
        # tg.append(ctm[i])
        # path = f'{out_dir}/{i}.TextGrid'
        # tg.write(path)
    p.close()
    p.join()
    for i in res:
        f.write(f'{i.get()[0]} {i.get()[1]}\n')
    f.close()


if __name__ == '__main__':
    # ctm_file = '/home/huangjunjie/workspace/Alignment/MultiAli/01_alignment_kaldi/v1/test/ctm/tri1_ali/ctm'
    # out_dir = '/home/huangjunjie/workspace/Alignment/MultiAli/01_alignment_kaldi/v1/test/ctm/tri1_ali/textgrid/'

    ctm_file = sys.argv[1]
    out_dir = sys.argv[2]
    ctm_to_textgrid(ctm_file, out_dir)
        